import pandas as pd
from PyQt5.QtWidgets import QMessageBox, QLabel, QComboBox, QLineEdit
import matplotlib.pyplot as plt
from PyQt5.QtCore import Qt
import numpy as np
import os

class RawDataProcessor:
    def __init__(self, parent):
        self.parent = parent
        
    def get_raw_data(self, contract_type):
        """获取原始数据，返回变化量(diff)形式"""
        # 获取时间范围
        start_date = self.parent.start_date_edit.text().strip()
        end_date = self.parent.end_date_edit.text().strip()
        
        # 验证日期格式
        if start_date and not self.is_valid_date(start_date):
            QMessageBox.warning(self.parent, "日期错误", "开始日期格式应为YYYYMMDD")
            return None, None
            
        if end_date and not self.is_valid_date(end_date):
            QMessageBox.warning(self.parent, "日期错误", "结束日期格式应为YYYYMMDD")
            return None, None
        
        try:
            # 读取净空单数据
            if not self.parent.file_paths[0]:
                QMessageBox.warning(self.parent, "文件错误", "净空单文件未选择")
                return None, None
            
            net_short_df = self.parent.read_file(self,self.parent.file_paths[0])
            
            # 验证净空单数据
            if contract_type == 'SH':
                # 对于上证指数，计算总净空单
                numeric_cols = [col for col in net_short_df.columns if col != '日期' and pd.api.types.is_numeric_dtype(net_short_df[col])]
                
                if not numeric_cols:
                    QMessageBox.warning(self.parent, "数据错误", "净空单文件中没有可用的合约列！")
                    return None, None
                
                # 计算总净空单
                net_short_df['总净空单'] = net_short_df[numeric_cols].sum(axis=1)
            else:
                # 对于其他合约，验证是否有对应列
                if contract_type not in net_short_df.columns:
                    QMessageBox.warning(self.parent, "数据错误", f"净空单文件中没有'{contract_type}'列！")
                    return None, None
            
            # 标准化日期格式
            net_short_df['日期'] = net_short_df['日期'].apply(lambda x: self.parent.normalize_date(self,x))
            # 过滤无效日期
            net_short_df = net_short_df[net_short_df['日期'] != ""]
            
            # 读取指数数据
            idx_map = {'IC': 1, 'IF': 2, 'IH': 3, 'IM': 4, 'SH': 5}
            idx = idx_map.get(contract_type)
            if idx is None or not self.parent.file_paths[idx]:
                return None, None
            
            index_df = self.parent.read_file(self,self.parent.file_paths[idx])
            
            # 验证指数数据
            required_columns = ['日期', '收盘价']
            for col in required_columns:
                if col not in index_df.columns:
                    # 尝试英文列名
                    col_map = {
                        '日期': ['date', 'Date', '时间'],
                        '收盘价': ['close', 'Close', '收盘']
                    }
                    
                    found = False
                    for alt in col_map.get(col, []):
                        if alt in index_df.columns:
                            index_df = index_df.rename(columns={alt: col})
                            found = True
                            break
                    
                    if not found:
                        QMessageBox.warning(self.parent, "数据错误", f"{contract_type}指数文件缺少'{col}'列！")
                        return None, None
            
            # 标准化日期格式
            index_df['日期'] = index_df['日期'].apply(lambda x: self.parent.normalize_date(self,x))
            # 过滤无效日期
            index_df = index_df[index_df['日期'] != ""]
            
            # 获取该合约的上市日期
            listing_date = self.parent.listing_dates[contract_type]
            
            # 复制净空单数据
            if contract_type == 'SH':
                net_short_data = net_short_df[['日期', '总净空单']].copy()
                net_short_data = net_short_data.rename(columns={'总净空单': '净空单'})
            else:
                net_short_data = net_short_df[['日期', contract_type]].copy()
                net_short_data = net_short_data.rename(columns={contract_type: '净空单'})
            
            # 过滤上市日期前的数据
            net_short_data = net_short_data[net_short_data['日期'].astype(int) >= listing_date]
            index_data = index_df[index_df['日期'].astype(int) >= listing_date]
            
            # 如果过滤后没有数据，返回空
            if net_short_data.empty or index_data.empty:
                return None, None
            
            # 计算变化量 (diff)
            net_short_data['净空单变化量'] = net_short_data['净空单'].diff()
            net_short_data['净空单变化百分比'] = net_short_data['净空单'].pct_change() * 100
            index_data = index_data.copy()  # 创建副本避免警告
            index_data.loc[:, '指数变化量'] = index_data['收盘价'].diff()
            index_data.loc[:, '指数变化百分比'] = index_data['收盘价'].pct_change() * 100
            
            # 过滤时间范围
            if start_date:
                net_short_data = net_short_data[net_short_data['日期'] >= start_date]
                index_data = index_data[index_data['日期'] >= start_date]
            
            if end_date:
                net_short_data = net_short_data[net_short_data['日期'] <= end_date]
                index_data = index_data[index_data['日期'] <= end_date]
            
            # 按日期排序
            net_short_data = net_short_data.sort_values('日期').reset_index(drop=True)
            index_data = index_data.sort_values('日期').reset_index(drop=True)
            
            return net_short_data, index_data
        
        except Exception as e:
            QMessageBox.critical(self.parent, "加载错误", f"加载原始数据时出错:\n{str(e)}")
            return None, None
            
    def is_valid_date(self, date_str):
        """验证日期格式是否为YYYYMMDD"""
        return len(date_str) == 8 and date_str.isdigit()
    
    def plot_raw_data(self, net_short_data, index_data, contract_type, display_type):
        """绘制原始数据图表（变化量形式）"""
        # 获取中文名称
        chinese_name = self.parent.get_chinese_name(contract_type)
        
        # 创建新的图表对象
        figure = plt.figure(figsize=(14, 10), dpi=100)
        
        # 设置支持中文的字体
        plt.rcParams['font.sans-serif'] = ['SimHei', 'Microsoft YaHei', 'WenQuanYi Zen Hei']
        plt.rcParams['axes.unicode_minus'] = False
        
        # 创建两个子图（上下排列，共享X轴）
        ax1 = figure.add_subplot(211)
        ax2 = figure.add_subplot(212, sharex=ax1)
        
        # 设置主标题
        start_date = net_short_data['日期'].min()
        end_date = net_short_data['日期'].max()
        figure.suptitle(f"{chinese_name}净空单变化量与指数变化量\n({start_date} 至 {end_date})", fontsize=14)
        
        # 绘制净空单变化量（上方子图）
        colors = ['g' if x >= 0 else 'r' for x in net_short_data['净空单变化量']]
        ax1.bar(range(len(net_short_data)), net_short_data['净空单变化量'], color=colors)
        ax1.axhline(y=0, color='k', linestyle='-', alpha=0.3)
        ax1.set_ylabel('净空单变化量', fontsize=12)
        
        # 绘制指数变化量（下方子图）
        colors = ['g' if x >= 0 else 'r' for x in index_data['指数变化量']]
        ax2.bar(range(len(index_data)), index_data['指数变化量'], color=colors)
        ax2.axhline(y=0, color='k', linestyle='-', alpha=0.3)
        ax2.set_ylabel('指数变化量', fontsize=12)
        
        # 设置X轴标签
        ax2.set_xlabel('时间序列', fontsize=12)
        
        # 设置X轴刻度格式
        date_labels = [str(int(date)) for date in net_short_data['日期']]
        
        # 如果日期太多，只显示部分标签
        if len(date_labels) > 20:
            step = max(1, len(date_labels) // 10)
            positions = list(range(0, len(date_labels), step))
            labels = date_labels[::step]
            ax2.set_xticks(positions)
            ax2.set_xticklabels(labels, rotation=45, ha='right')
        else:
            ax2.set_xticks(range(len(net_short_data)))
            ax2.set_xticklabels(date_labels, rotation=45, ha='right')
        
        # 添加网格
        ax1.grid(True, linestyle='--', alpha=0.6)
        ax2.grid(True, linestyle='--', alpha=0.6)
        
        # 调整布局
        figure.tight_layout()
        figure.subplots_adjust(top=0.9)  # 为总标题留出空间
        
        return figure
    
    def export_raw_data(self, net_short_data, index_data, contract_type, output_dir):
        """导出原始数据到CSV（变化量形式）"""
        # 获取合约中文名称
        chinese_name = self.parent.get_chinese_name(contract_type)
        
        # 生成文件名
        start_date = net_short_data['日期'].min()
        end_date = net_short_data['日期'].max()
        
        # 保存净空单变化量数据
        net_short_path = os.path.join(output_dir, f"{chinese_name}_净空单变化量_{start_date}_{end_date}.csv")
        net_short_data.to_csv(net_short_path, index=False, float_format='%.4f')
        
        # 保存指数变化量数据
        index_path = os.path.join(output_dir, f"{chinese_name}_指数变化量_{start_date}_{end_date}.csv")
        index_data.to_csv(index_path, index=False, float_format='%.4f')
        
        return net_short_path, index_path